import os, re, sys, copy
import json5
import numpy as np
import time
import random
import sympy as sym
from utils import json_tools
try:
    from transformations import identity_matrix, transform_points, concatenate_matrices, rotation_matrix, translation_matrix
except:
    from .transformations import *

try:
    # from part1 import Part, PART_DIR, GROUP_DIR,Global_Cache
    from jxq.part import *
except:
    # from .part1 import Part, PART_DIR, GROUP_DIR,Global_Cache
    from .part import *

__version__ = "1.7.0"
print('AIAD 3D Parser', __version__)


def get_extra_args(args, k):
    arr = []
    for name in args:
        ss = name.split('__', 1)
        if ss[0] == k:
            dic = {}
            dic.update(args[name])
            dic['name'] = ss[1]

            arr.append(dic)

    return arr


class Group(Part):
    def __init__(self, group_type,
                 group_dir=None,
                 part_dir=None,json={},use_json=False, extra_args=None, global_cache=None
                 ):
        self.use_json=use_json
        if global_cache is not None:
            Global_Cache.Cache = copy.deepcopy(global_cache)
        self.parse_element(group_type, group_dir)

        self.parse_extra_args(extra_args)

        for k in self.elements:
            element = self.elements[k]
            category = element['category']

            sub_extra_args = get_extra_args(self.args, k)

            if category.startswith('.'):
                element['instance'] = Group(category[1:], group_dir, part_dir, extra_args=sub_extra_args)
            else:
                # if element['desc']=="step":
                #     element['instance']=Part(category+"_stp",part_dir, extra_args=sub_extra_args)
                # else:
                element['instance'] = Part(category, part_dir, extra_args=sub_extra_args)

        self.get_transform()

    def get_child(self, name):
        ss = name.split('/')
        if len(ss) == 1:
            return self.elements[name]['instance'], self.elements[name]['transform']
        else:
            raise NotImplementedError()

    def get_point(self, name):
        if isinstance(name, list):
            assert len(name) == 3
            return np.array(name, dtype=float)

        if '/' not in name:
            assert name in self.points

            refer = self.points[name].get('refer')
            offset = self.points[name].get('offset')

            if refer is None:
                coor = np.array([0, 0, 0], np.float32)
            else:
                assert '/' in refer

                coor = Group.get_point(self, refer)

            if offset is not None:
                coor += np.array(offset)
        else:
            ss = name.split('/')
            assert len(ss) > 1
            # element = self
            m = identity_matrix()
            for e in ss[:-1]:
                element, tran = self.get_child(e)
                m = concatenate_matrices(m, tran)
            # print(m)

            coor = element.get_point(ss[-1])
            coor = transform_points([coor], m)[0]

        return coor

    def shape(self, debug=False):
        mesh_all = None
        name_all =[]
        for name in self.elements:
            name_all.append(name)
            element = self.elements[name]
            m = element['instance'].compound()
            a = [[element['transform'][i][j] for j in range(4)] for i in range(3)]

            T = gp_Trsf()
            myshape = TopoDS_Shape(m)
            T.SetValues(a[0][0], a[0][1], a[0][2], a[0][3],
                        a[1][0], a[1][1], a[1][2], a[1][3],
                        a[2][0], a[2][1], a[2][2], a[2][3])
            trans = BRepBuilderAPI_Transform(T)
            trans.Perform(myshape)
            new_shape = trans.Shape()
            if mesh_all is None:
                mesh_all = [new_shape]
            else:
                mesh_all.append(new_shape)
        return mesh_all,name_all

    def compound(self):
        myNewShape = self.shape()
        if isinstance(myNewShape,list):
            compound = TopoDS_Compound()
            B = BRep_Builder()
            B.MakeCompound(compound)
            for shape in myNewShape:
                B.Add(compound, shape)

            return compound

        else:
            compound = TopoDS_Compound()
            B = BRep_Builder()
            B.MakeCompound(compound)
            B.Add(compound,myNewShape)
            return compound

    def to_stl(self,save_dir,filename):
        # make sure the path exists otherwise OCE get confused
        if not os.path.isdir(save_dir):
            raise AssertionError("wrong path provided")
        stl_file = os.path.join(save_dir, filename)
        write_stl_file(self.compound(), stl_file)


class Body(Group):
    def __init__(self, group_type,
                 group_dir=None,
                 part_dir=None, extra_args=None,
                 transform=None, global_cache=None
                 ):
        self.path = os.path.join(group_dir, group_type)
        self.group_type = group_type
        self.group_dir = group_dir
        self.part_dir = part_dir

        if global_cache is not None:
            Global_Cache.Cache = copy.deepcopy(global_cache)
        self.transform = transform

        if self.transform is None:
            self.transform = identity_matrix()

        super().__init__(group_type, group_dir=group_dir, part_dir=part_dir, extra_args=extra_args,
                         global_cache=global_cache)

    def apply_transform(self, mat):
        self.transform = concatenate_matrices(mat, self.transform)

    def add_transform_sym(self, group, m):
        assert isinstance(group, list)

        for e in group:
            element = self.elements[e]
            element['transform_sym'] = m * element['transform_sym']

    def get_transform_m(self):
        return copy.deepcopy(self.transform)

    def get_point(self, name):
        coor = super().get_point(name)
        coor = transform_points([coor], self.transform)[0]
        return coor

    def get_elements_name_list(self):
        return list(self.elements.keys())

    def get_object_list(self):
        for name in self.elements:
            element = self.elements[name]
            element['transform'] = concatenate_matrices(
                self.transform, element['transform'])

    def apply_translation(self, v):
        self.transform = concatenate_matrices(translation_matrix(v), self.transform)

    def move(self, p):
        self.transform = concatenate_matrices(translation_matrix(p), self.transform)
        random_num = random.randrange(100, 65535)
        self.trans[str(random_num)] = {}
        self.trans[str(random_num)]['group'] = self.get_elements_name_list()
        self.trans[str(random_num)]['move'] = p

    def move_p2p(self, p1, p2):
        p = p2 - p1
        self.transform = concatenate_matrices(translation_matrix(p), self.transform)
        random_num = random.randrange(100, 65535)
        self.trans[str(random_num)] = {}
        self.trans[str(random_num)]['group'] = self.get_elements_name_list()
        self.trans[str(random_num)]['point2point'] = [p1.tolist(), p2]

    def rotate(self, axis, angle):
        random_num = random.randrange(65535, 655355)
        self.trans[str(random_num)] = {}
        self.trans[str(random_num)]['group'] = self.get_elements_name_list()
        self.trans[str(random_num)]['rotate'] = [angle, axis]
        json_tools.save_json(self.to_json(), self.path + '_1.json5')
        return Body(self.group_type + '_1', self.group_dir, self.part_dir)

    def shape(self, debug=False):
        mesh_all = super().shape(debug)[0]
        myNewShape = mesh_all
        compound = TopoDS_Compound()
        B = BRep_Builder()
        B.MakeCompound(compound)
        for shape in myNewShape:
            B.Add(compound, shape)
        a = [[self.transform[i][j] for j in range(4)] for i in range(3)]

        T = gp_Trsf()
        myshape = TopoDS_Shape(compound)
        T.SetValues(a[0][0], a[0][1], a[0][2], a[0][3],
                    a[1][0], a[1][1], a[1][2], a[1][3],
                    a[2][0], a[2][1], a[2][2], a[2][3])
        trans = BRepBuilderAPI_Transform(T)
        trans.Perform(myshape)
        new_shape = trans.Shape()

        return new_shape

    def shape_dict(self,debug=False):
        mesh_all,name_all = super().shape(debug)
        print(len(mesh_all))
        # myNewShape = mesh_all
        shape_dict={}

        a = [[self.transform[i][j] for j in range(4)] for i in range(3)]

        T = gp_Trsf()
        T.SetValues(a[0][0], a[0][1], a[0][2], a[0][3],
                    a[1][0], a[1][1], a[1][2], a[1][3],
                    a[2][0], a[2][1], a[2][2], a[2][3])
        trans = BRepBuilderAPI_Transform(T)
        for i in range(len(name_all)):
            trans.Perform(mesh_all[i])
            new_shape = trans.Shape()
            shape_dict[name_all[i]]=new_shape
        # for shape in myNewShape:
        #     shape_dict[""]=trans.Perform(shape)

        return shape_dict

    def to_stl(self,save_dir,filename):
        # make sure the path exists otherwise OCE get confused
        if not os.path.isdir(save_dir):
            raise AssertionError("wrong path provided")
        stl_file = os.path.join(save_dir, filename)
        write_stl_file(self.shape(), stl_file)

    def to_json(self):
        ans = {}
        ans['version'] = 4.0
        ans['elements'] = []
        ans['trans'] = []
        for ele in self.elements.keys():
            ele_child = {}
            ele_child['name'] = ele
            ele_child['category'] = self.elements[ele]['category']
            ele_child['desc'] = self.elements[ele]['desc']
            ans['elements'].append(ele_child)
        for tran in self.trans.keys():
            tran_child = {'group': self.trans[tran]['group']}
            if 'point2point' in self.trans[tran].keys():
                tran_child['point2point'] = self.trans[tran]['point2point']
            if 'move' in self.trans[tran].keys():
                tran_child['move'] = self.trans[tran]['move']
            if 'rotate' in self.trans[tran].keys():
                tran_child['rotate'] = self.trans[tran]['rotate']
            ans['trans'].append(tran_child)
            
        # random_num = random.randrange(0, 65535)
        # temp_ans = json5.dumps(json5.loads(ans), ensure_ascii=False, indent=4)
        # for ele in self.elements.keys():
        #     temp_ans = temp_ans.replace(ele, ele + '_' + str(random_num))
        # ans = json5.loads(temp_ans)

        # for ele in ans['elements']:
        #     ele['name'] += '_' + str(random_num)
        # for tran in ans['trans']:
        #     for key in tran.keys():
        #         if key == 'group':
        #             for i in range(len(tran[key])):
        #                 tran[key][i] += '_' + str(random_num)
        #         if key == 'point2point':
        #             for i in range(len(tran[key])):
        #                 tran[key][i] = tran[key][i].split('/')[0] + '_' + str(random_num) + '/' + tran[key][i].split('/')[1]
        return ans


if __name__ == '__main__':
    import time
    from OCC.Display.SimpleGui import init_display

    display, start_display, add_menu, add_function_to_menu = init_display("wx")

    start = time.perf_counter()  # 起始时间
    name = 'group'
    group = Body(name,group_dir=r'C:\Users\Lenovo\Desktop\5月用\29-02 - 副本',part_dir=r'C:\Users\Lenovo\Desktop\5月用\29-02 - 副本')
    print(group.args)
    display.DisplayShape(group.compound(), update=True)
    start_display()




